import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from diffusers import StableVideoDiffusionPipeline
from PIL import Image
import requests
from io import BytesIO
from moviepy import ImageSequenceClip
import argparse
import psutil
import socket
import netifaces
from datetime import timedelta


def print_resource_usage(rank):
    print(f"\n【节点{rank}】资源监控：")
    if torch.cuda.is_available():
        print(f"GPU{rank} 内存 - 已分配: {torch.cuda.memory_allocated() / 1024 ** 2:.2f}MB")
    print(f"CPU使用率: {psutil.cpu_percent()}%")


def get_physical_ips():
    """获取物理网卡的真实IPv4地址"""
    physical_ips = []
    for iface in netifaces.interfaces():
        # 排除虚拟网卡（匹配中文和英文系统）
        if any(keyword in iface for keyword in ['Virtual', 'VMware', 'VPN', 'Loopback']):
            continue
        addrs = netifaces.ifaddresses(iface)
        if netifaces.AF_INET in addrs:
            for addr in addrs[netifaces.AF_INET]:
                ip = addr['addr']
                # 排除本地回环和链路本地地址
                if not ip.startswith('127.') and not ip.startswith('169.254.'):
                    physical_ips.append(ip)
    return physical_ips


def setup(rank, world_size, master_ip):
    """初始化分布式环境（关键修正部分）"""
    # 验证主节点IP有效性
    physical_ips = get_physical_ips()
    if master_ip not in physical_ips:
        raise ValueError(f"主节点IP {master_ip} 无效，可用物理IP列表：{physical_ips}")

    # 配置Gloo网络参数（Windows专用设置）
    os.environ['MASTER_ADDR'] = '192.168.12.251'
    os.environ['MASTER_PORT'] = '8080'
    os.environ['GLOO_SOCKET_IFNAME'] = '以太网' if '以太网' in netifaces.interfaces() else 'Ethernet'

    # 初始化进程组（增加超时设置）
    try:
        dist.init_process_group(
            backend="gloo",
            rank=rank,
            world_size=world_size,
            timeout=timedelta(seconds=30)
        )
        print(f"【节点{rank}】成功连接到主节点 {master_ip}:29500")
    except RuntimeError as e:
        print(f"【节点{rank}】连接失败：{str(e)}")
        raise


def cleanup():
    """安全清理分布式环境"""
    if dist.is_initialized():
        dist.destroy_process_group()
    torch.cuda.empty_cache()


def main(rank, world_size, master_ip):
    try:
        # 初始化分布式环境
        setup(rank, world_size, master_ip)

        # 加载模型到指定GPU
        torch.cuda.set_device(rank)
        pipeline = StableVideoDiffusionPipeline.from_pretrained(
            "stabilityai/stable-video-diffusion-img2vid-xt",
            torch_dtype=torch.float32,
            cache_dir='models'
        ).to(f"cuda:{rank}")
        ddp_pipeline = DDP(pipeline, device_ids=[rank])

        # 主节点加载数据
        if rank == 0:
            url = "https://pic.rmb.bdstatic.com/bjh/other/7132183c3a3973c03796a833e857df67.jpeg"
            response = requests.get(url, timeout=10)
            init_image = Image.open(BytesIO(response.content)).convert("RGB")
        else:
            init_image = None

        # 广播图像数据（转换为Tensor确保兼容性）
        if rank == 0:
            tensor_image = pipeline.image_processor.preprocess(init_image)
        else:
            tensor_image = torch.zeros_like(pipeline.image_processor.preprocess(Image.new('RGB', (512, 512))))

        dist.broadcast(tensor_image, src=0)
        print_resource_usage(rank)

        # 执行推理
        with torch.no_grad():
            video_frames = ddp_pipeline(tensor_image, num_frames=8).frames

        # 主节点保存结果
        if rank == 0:
            clip = ImageSequenceClip([frame for frame in video_frames], fps=4)
            os.makedirs("output", exist_ok=True)
            clip.write_videofile("output/video.mp4", codec="libx264", verbose=False)
            print("视频生成成功")

    except Exception as e:
        print(f"【节点{rank}】发生错误：{str(e)}")
    finally:
        cleanup()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--rank", type=int, required=True)
    parser.add_argument("--world_size", type=int, required=True)
    parser.add_argument("--master_addr", type=str, required=True)
    args = parser.parse_args()

    main(args.rank, args.world_size, args.master_addr)
